# -*- coding: utf-8 -*-

import os
import cv2
import math
import time
import numpy as np
from utils import stretch_linear, calc_cc3d, put_heatmap_3d


scan_dir = "Pneumonia_sample_corrected/others/FE923712_les"  # mycoplasma/FE781707_les##
dst_dir = "Results2/others/FE923712_les"  # mycoplasma/FE781707_les##
scan_lst = sorted(os.listdir(scan_dir))
img_lst = []
for scan in scan_lst:
    scan_arr = cv2.imread(os.path.join(scan_dir, scan), 0)
    img_lst.append(scan_arr)
img_arr = np.array(img_lst, dtype=np.int64) // 255
start_t = time.time()
ccinfo_dct = calc_cc3d(img_arr)  # 一个Volume获取3D连通域信息
end_t = time.time()
print("Time cost: %f s. " % (end_t - start_t))

scov_arr = 4.0 * np.eye(3)
heatmap = np.zeros_like(img_arr, np.float64)
start_t = time.time()
for idx in range(len(ccinfo_dct["S"])):  # 每个3D连通域构建3D高斯热力图
    # 若3D连通域为单像素或其协方差为标量,则指定其协方差矩阵为固定大小对角线元素的对角矩阵构建3D高斯热力图
    # 否则按其协方差矩阵构建3D高斯热力图
    if ccinfo_dct["S"][idx] == 1 or ccinfo_dct["cov"][idx].shape != (3, 3):
        heatmap = put_heatmap_3d(heatmap, ccinfo_dct["mean"][idx], scov_arr)
    else:
        heatmap = put_heatmap_3d(heatmap, ccinfo_dct["mean"][idx], ccinfo_dct["cov"][idx])
end_t = time.time()
print("Time cost: %f s. " % (end_t - start_t))

heatmap = stretch_linear(heatmap)
os.makedirs(dst_dir, exist_ok=True)
for idx in range(heatmap.shape[0]):
    cv2.imwrite(os.path.join(dst_dir, "%03d.png"%(idx+1)), heatmap[idx, :, :])
